# encoding: utf-8
import openai
import yaml
import argparse

class Ai00:
    def __init__(self, model="model", port=65530, api_key="permisionkey", prompt_key="default", max_history=5):
        self.initialize_openai(port, api_key)
        self.ctx = []
        self.params = self.default_params(model)
        self.load_prompt(prompt_key)
        self.max_history = max_history
        assistant_name = self.params["assistant_name"]
        welcome_message = self.params.get("welcome_message", f"我是{assistant_name}，很开心和你聊天～")
        self.append_ctx("assistant", welcome_message)
        print(f"{assistant_name}：{welcome_message}")

    def initialize_openai(self, port, api_key):
        openai.api_base = f"http://127.0.0.1:{port}/api/oai"
        openai.api_key = api_key

    def default_params(self, model):
        return {
            "system_name": "System",
            "user_name": "User",
            "assistant_name": "Assistant",
            "model": model,
            "max_tokens": 4096,
            "top_p": 0.6,
            "temperature": 1,
            "presence_penalty": 0.3,
            "frequency_penalty": 0.3,
            "half_life": 400,
            "stop": ['\x00', '\n\n']
        }

    def load_prompt(self, prompt_key):
        with open("template/prompts.yaml", "r", encoding="utf-8") as file:
            prompts = yaml.safe_load(file)
            prompt_data = prompts["prompts"].get(prompt_key, {})
            self.params["user_name"] = prompt_data.get("user_name", "User")
            self.params["assistant_name"] = prompt_data.get("assistant_name", "Assistant")
            self.params["welcome_message"] = prompt_data.get("welcome_message", f"我是{self.params['assistant_name']}，很开心和你聊天～")
            self.append_ctx("system", prompt_data.get("prompt", ""))

    def set_params(self, **kwargs):
        self.params.update(kwargs)

    def clear_ctx(self, save_system=False):
        self.ctx = self.ctx[:1] if save_system else []

    def get_ctx(self):
        return self.ctx

    def append_ctx(self, role, content):
        self.ctx.append({"role": role, "content": content})

    def send_message(self, message, save_ctx=True, role="user"):
        self.append_ctx(role, message)
        # 如果历史记录超过最大次数，从1开始删除最早的记录
        # 保留第0个记录，即系统提示
        if len(self.ctx) > self.max_history:
            self.ctx.pop(1)

        result = openai.ChatCompletion.create(
            model=self.params['model'],
            messages=self.ctx,
            names={
                "system": self.params['system_name'],
                "user": self.params['user_name'],
                "assistant": self.params['assistant_name']
            },
            max_tokens=self.params['max_tokens'],
            half_life=self.params['half_life'],
            top_p=self.params['top_p'],
            temperature=self.params['temperature'],
            presence_penalty=self.params['presence_penalty'],
            frequency_penalty=self.params['frequency_penalty'],
            stop=self.params['stop']
        ).choices[0].message['content']

        if save_ctx:
            self.append_ctx("assistant", result)
        return result

    def new_conversation(self):
        self.clear_ctx(save_system=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run Ai00 with specific prompt")
    parser.add_argument("--prompt", type=str, default="default", help="角色代码")
    parser.add_argument("--max_history", type=int, default=7, help="对话时引用的历史记录次数")
    args = parser.parse_args()

    print("### 输入 `/new` 开始新对话，输入 `/exit` 退出程序 ###")
    print("-----------------------------------------------")
    print()

    ai = Ai00(prompt_key=args.prompt, max_history=args.max_history)
    ai.set_params(
        top_p=0.8,
        temperature=1.6,
        presence_penalty=0.3,
        frequency_penalty=0.5,
        stops=['\x00', '\n'],
    )

    user_name = ai.params["user_name"]
    assistant_name = ai.params["assistant_name"]

    try:
        while True:
            message = input(f"{user_name}：")
            if message == "/exit":
                break
            elif message == "/new":
                ai.new_conversation()
                print(f"{assistant_name}：一个崭新的{assistant_name}回来了～")
            else:
                response = ai.send_message(message)
                print(f"{assistant_name}：{response}")
                print()
    except KeyboardInterrupt:
        print("\n程序已终止")